import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import adfuller, acf
from scipy import stats
import talib as ta
from collections import deque
import time

class MarketRegimeDetector:
    """市场状态检测器，用于识别趋势、震荡和极端行情"""
    def __init__(self, config=None):
        self.config = config or {
            'lookback_period': 60,  # 回看周期（分钟）
            'volatility_threshold': 2.0,  # 波动率阈值倍数
            'trend_strength_threshold': 0.5,  # 趋势强度阈值
            'adf_threshold': 0.05,  # ADF检验显著性水平
            'regime_change_threshold': 0.3  # 状态变化阈值
        }
        
        # 历史状态记录
        self.regime_history = deque(maxlen=1000)
        self.volatility_history = deque(maxlen=200)  # 存储波动率历史
        self.trend_strength_history = deque(maxlen=200)  # 存储趋势强度历史
        
        # 当前市场状态
        self.current_regime = 'NEUTRAL'
        self.last_regime = 'NEUTRAL'
        self.regime_confidence = 0.0
        
        # 市场状态常量
        self.REGIMES = {
            'TREND_BULLISH': '上升趋势',
            'TREND_BEARISH': '下降趋势',
            'RANGE_BOUND': '震荡行情',
            'EXTREME_VOLATILITY': '极端波动',
            'NEUTRAL': '中性市场'
        }
    
    def calculate_volatility(self, prices):
        """计算波动率"""
        if len(prices) < 2:
            return 0.0
        
        # 计算对数收益率
        returns = np.log(prices[1:] / prices[:-1])
        
        # 计算年化波动率（假设一天240分钟）
        volatility = np.std(returns) * np.sqrt(240 * 252)  # 年化
        
        # 记录波动率历史
        self.volatility_history.append(volatility)
        
        return volatility
    
    def calculate_trend_strength(self, prices):
        """计算趋势强度"""
        if len(prices) < 20:
            return 0.0
        
        # 使用ADF检验判断是否存在趋势
        adf_result = adfuller(prices)
        adf_pvalue = adf_result[1]
        
        # 计算MACD指标
        macd, macd_signal, macd_hist = ta.MACD(
            np.array(prices),
            fastperiod=12, slowperiod=26, signalperiod=9
        )
        
        # 计算RSI指标
        rsi = ta.RSI(np.array(prices), timeperiod=14)
        
        # 计算布林带
        upper, middle, lower = ta.BBANDS(np.array(prices), timeperiod=20)
        
        # 计算趋势强度分数（综合多个指标）
        # 1. ADF检验结果（趋势存在的概率）
        adf_score = 1.0 if adf_pvalue > self.config['adf_threshold'] else 0.0
        
        # 2. MACD交叉信号
        if len(macd_hist) > 0 and not np.isnan(macd_hist[-1]):
            macd_score = 0.5 * (1 + np.tanh(macd_hist[-1] / np.std(macd_hist[~np.isnan(macd_hist)]))) if len(macd_hist[~np.isnan(macd_hist)]) > 0 else 0.5
        else:
            macd_score = 0.5
        
        # 3. 价格在布林带中的位置
        if len(prices) > 0 and len(upper) > 0 and not np.isnan(upper[-1]) and not np.isnan(lower[-1]) and upper[-1] != lower[-1]:
            bollinger_score = (prices[-1] - lower[-1]) / (upper[-1] - lower[-1])
        else:
            bollinger_score = 0.5
        
        # 4. RSI指标
        if len(rsi) > 0 and not np.isnan(rsi[-1]):
            rsi_score = (rsi[-1] - 30) / 40  # 映射到0-1区间
            rsi_score = max(0, min(1, rsi_score))  # 限制范围
        else:
            rsi_score = 0.5
        
        # 综合评分（可以根据实际效果调整权重）
        trend_strength = 0.3 * adf_score + 0.2 * macd_score + 0.2 * bollinger_score + 0.3 * rsi_score
        
        # 记录趋势强度历史
        self.trend_strength_history.append(trend_strength)
        
        return trend_strength
    
    def detect_market_regime(self, prices, volumes=None):
        """检测当前市场状态"""
        if len(prices) < self.config['lookback_period']:
            return {
                'regime': 'NEUTRAL',
                'confidence': 0.0,
                'message': 'Insufficient data for regime detection'
            }
        
        # 计算关键指标
        volatility = self.calculate_volatility(prices)
        trend_strength = self.calculate_trend_strength(prices)
        
        # 检查是否为极端波动行情
        is_extreme_volatility = self._is_extreme_volatility(volatility, volumes)
        
        if is_extreme_volatility:
            regime = 'EXTREME_VOLATILITY'
            confidence = 0.9  # 极端波动的置信度通常较高
        else:
            # 判断是趋势还是震荡行情
            if trend_strength > self.config['trend_strength_threshold']:
                # 上升趋势
                regime = 'TREND_BULLISH'
                confidence = min(1.0, trend_strength)
            elif trend_strength < (1.0 - self.config['trend_strength_threshold']):
                # 下降趋势
                regime = 'TREND_BEARISH'
                confidence = min(1.0, 1.0 - trend_strength)
            else:
                # 震荡行情
                regime = 'RANGE_BOUND'
                confidence = min(1.0, 1.0 - abs(0.5 - trend_strength) * 2)
        
        # 更新当前状态
        self.last_regime = self.current_regime
        self.current_regime = regime
        self.regime_confidence = confidence
        
        # 记录状态历史
        current_time = int(time.time() * 1000)
        self.regime_history.append({
            'timestamp': current_time,
            'regime': regime,
            'confidence': confidence,
            'volatility': volatility,
            'trend_strength': trend_strength
        })
        
        # 检查状态是否发生变化
        regime_changed = self._check_regime_change()
        
        return {
            'regime': regime,
            'regime_name': self.REGIMES[regime],
            'confidence': confidence,
            'volatility': volatility,
            'trend_strength': trend_strength,
            'regime_changed': regime_changed,
            'timestamp': current_time
        }
    
    def _is_extreme_volatility(self, current_volatility, volumes=None):
        """判断是否为极端波动行情"""
        # 如果波动率历史不足，使用绝对阈值判断
        if len(self.volatility_history) < 20:
            return current_volatility > self.config['volatility_threshold']
        
        # 计算历史波动率的均值和标准差
        historical_vols = np.array(self.volatility_history)
        mean_vol = np.mean(historical_vols)
        std_vol = np.std(historical_vols)
        
        # 计算当前波动率相对于历史的z-score
        z_score = (current_volatility - mean_vol) / std_vol if std_vol > 0 else 0
        
        # 判断是否为极端波动（例如超过2个标准差）
        is_extreme = abs(z_score) > self.config['volatility_threshold']
        
        # 如果有成交量数据，结合成交量判断
        if volumes is not None and len(volumes) > 20:
            # 计算成交量的z-score
            mean_volume = np.mean(volumes[-20:])
            std_volume = np.std(volumes[-20:])
            volume_z_score = (volumes[-1] - mean_volume) / std_volume if std_volume > 0 else 0
            
            # 成交量异常放大且波动率异常，更可能是极端行情
            if volume_z_score > 2.0 and abs(z_score) > 1.5:
                is_extreme = True
        
        return is_extreme
    
    def _check_regime_change(self):
        """检查市场状态是否发生变化"""
        if self.last_regime == self.current_regime:
            return False
        
        # 检查状态变化的置信度
        if self.regime_confidence < self.config['regime_change_threshold']:
            return False
        
        return True
    
    def get_regime_statistics(self, lookback=100):
        """获取市场状态统计信息"""
        if not self.regime_history:
            return {
                'total_samples': 0,
                'regime_distribution': {},
                'average_confidence': 0.0
            }
        
        # 获取最近lookback个状态记录
        recent_history = list(self.regime_history)[-lookback:]
        
        # 计算状态分布
        regime_counts = {}
        total_confidence = 0.0
        
        for record in recent_history:
            regime = record['regime']
            regime_counts[regime] = regime_counts.get(regime, 0) + 1
            total_confidence += record['confidence']
        
        # 计算状态概率
        regime_distribution = {}
        total = len(recent_history)
        
        for regime, count in regime_counts.items():
            regime_distribution[regime] = {
                'count': count,
                'probability': count / total,
                'regime_name': self.REGIMES[regime]
            }
        
        return {
            'total_samples': total,
            'regime_distribution': regime_distribution,
            'average_confidence': total_confidence / total if total > 0 else 0.0
        }
    
    def get_regime_signals(self, prices, volumes=None):
        """获取基于市场状态的交易信号"""
        regime_info = self.detect_market_regime(prices, volumes)
        regime = regime_info['regime']
        
        # 根据市场状态生成不同的交易信号建议
        signals = []
        
        if regime == 'TREND_BULLISH':
            # 上升趋势：建议跟随趋势
            signals.append({
                'signal_type': 'follow_trend',
                'direction': 'long',
                'strength': min(1.0, regime_info['confidence'] * 2),
                'reason': f'市场处于上升趋势，置信度{regime_info['confidence']:.2f}'
            })
        elif regime == 'TREND_BEARISH':
            # 下降趋势：建议做空或观望
            signals.append({
                'signal_type': 'follow_trend',
                'direction': 'short',
                'strength': min(1.0, regime_info['confidence'] * 2),
                'reason': f'市场处于下降趋势，置信度{regime_info['confidence']:.2f}'
            })
        elif regime == 'RANGE_BOUND':
            # 震荡行情：建议高抛低吸
            if len(prices) > 20:
                # 计算布林带
                upper, middle, lower = ta.BBANDS(np.array(prices), timeperiod=20)
                if len(upper) > 0 and not np.isnan(upper[-1]) and not np.isnan(lower[-1]):
                    current_price = prices[-1]
                    # 接近下轨，考虑买入
                    if current_price < lower[-1] * 1.01:
                        signals.append({
                            'signal_type': 'range_bound',
                            'direction': 'long',
                            'strength': 0.7,
                            'reason': '价格接近布林带下轨，震荡行情中可能反弹'
                        })
                    # 接近上轨，考虑卖出
                    elif current_price > upper[-1] * 0.99:
                        signals.append({
                            'signal_type': 'range_bound',
                            'direction': 'short',
                            'strength': 0.7,
                            'reason': '价格接近布林带上轨，震荡行情中可能回落'
                        })
        elif regime == 'EXTREME_VOLATILITY':
            # 极端波动：建议观望或轻仓
            signals.append({
                'signal_type': 'risk_management',
                'action': 'reduce_position',
                'strength': 1.0,
                'reason': '市场处于极端波动状态，建议控制风险'
            })
            signals.append({
                'signal_type': 'risk_management',
                'action': 'wait_for_clarity',
                'strength': 0.9,
                'reason': '市场波动剧烈，建议等待趋势明确后再入场'
            })
        
        # 添加风险提示
        risk_level = self._assess_risk_level(regime_info)
        
        return {
            'regime_info': regime_info,
            'signals': signals,
            'risk_level': risk_level,
            'timestamp': int(time.time() * 1000)
        }
    
    def _assess_risk_level(self, regime_info):
        """评估当前市场风险水平"""
        regime = regime_info['regime']
        volatility = regime_info['volatility']
        
        # 基础风险分数
        base_risk = 0.5
        
        # 根据市场状态调整风险分数
        if regime == 'EXTREME_VOLATILITY':
            base_risk = 0.9
        elif regime == 'TREND_BULLISH' or regime == 'TREND_BEARISH':
            base_risk = 0.6
        elif regime == 'RANGE_BOUND':
            base_risk = 0.4
        
        # 根据波动率调整风险分数
        if len(self.volatility_history) > 20:
            historical_vols = np.array(self.volatility_history)
            mean_vol = np.mean(historical_vols)
            
            # 波动率越高，风险越大
            vol_factor = min(1.0, volatility / (mean_vol * 2))
            base_risk = base_risk * (1 + vol_factor)
        
        # 确保风险分数在0-1之间
        risk_score = min(1.0, max(0.1, base_risk))
        
        # 映射到风险等级
        if risk_score > 0.8:
            risk_level = 'VERY_HIGH'
        elif risk_score > 0.6:
            risk_level = 'HIGH'
        elif risk_score > 0.4:
            risk_level = 'MEDIUM'
        else:
            risk_level = 'LOW'
        
        return {
            'risk_score': risk_score,
            'risk_level': risk_level,
            'factors': {
                'regime': regime,
                'volatility': volatility
            }
        }